#include <bcc/BPF.h>
#include <iostream>
#include <fstream>
#include <csignal>

//控制Ctrl^C信号
namespace TerSignal
{
    static bool keepRunning = false;
    void sig_handler(int sig){
        if (sig == SIGINT || sig==SIGALRM){
            keepRunning = false;
        }
    }
    void salarm(unsigned int n){
        alarm(n);
    }
    void init(){
        keepRunning = true;
        signal(SIGINT, sig_handler);
        signal(SIGALRM,sig_handler);
    }
    bool ok(){
        return keepRunning;
    }
}

#define BPF_FUN_START "fun_start"
#define BPF_FUN_END "fun_end"
#define BPF_FILE_NAME "tfun.c"
#define BPF_OUTPUT_NAME "output"

using namespace std;

typedef unsigned long long u64;
struct event_t{
    u64 time_delta;
};

class tfun{
    std::string fun;//要跟踪的函数名
    ebpf::BPF *bpf;

public:
    tfun(string fun);
    ~tfun();
    bool load();        //初始化bpf
    bool attach();      //挂上目标函数
    bool run(unsigned int timenum=0,string file="");         //运行跟踪过程
    //打印每次的跟踪结果
    static void print_event(void *cb_cookie, void *data, int data_size);
};

#define TFUN_ARGS "ht:f:o:"
const string help_str="usage: tfun.py [-h] [-f FNAME] [-t TIME] [-o OUTPUT]\n\
\n\
打印目标内核函数的执行时间\n\
\n\
optional arguments:\n\
  -h, --help            show this help message and exit\n\
  -f FNAME, --fname FNAME  要跟踪的内核函数名\n\
  -t TIME, --time TIME  程序要执行的时间\n\
  -o OUTPUT, --output OUTPUT  将结果写入目标文件";

int main(int argc, char* argv[])
{
    int arg;
    string fun="";
    string output_file;
    unsigned int time;
    while((arg=getopt(argc,argv,TFUN_ARGS))!=-1){
        switch(arg){
            case 'h':{
                cout<<help_str<<endl;
                break;
            }
            case 'f':
                fun=optarg;
                break;
            case 't':
                time=atol(optarg);
                break;
            case 'o':
                output_file=optarg;
                break;
            case '?':
                break;
        }
    }
    if (fun==""){
        return 0;
    }
    tfun tf(fun);
    if( !tf.load())
        return 1;
    if (!tf.attach())
        return 1;
    if (!tf.run(time,output_file))
        return 1;
    return 0; 
}

tfun::tfun(string fun)
{
    this->fun=fun;
}

tfun::~tfun()
{
}

bool tfun::load()
{
    ifstream fin(BPF_FILE_NAME,ios::binary);
    if (!fin.is_open()){
        return false;
    }
    istreambuf_iterator<char>beg(fin),end;
    string ebpf_file_str(beg,end);
    this->bpf=new ebpf::BPF();
    if (this->bpf==nullptr)
        return false;
    this->bpf->init(ebpf_file_str);
    return true;
}

bool tfun::attach()
{
    if (!this->bpf->attach_kprobe(this->fun,BPF_FUN_START).ok())
        return false;
    if (!this->bpf->attach_kprobe(this->fun,BPF_FUN_END,0,BPF_PROBE_RETURN).ok())
            return false;
    return true;
}

//@cb_cookie: 自动填充的cpu数据
//@data: ebpf代码中perf_submit传出的结构
//@data_size: data结构大小
void tfun::print_event(void *cb_cookie, void *data, int data_size)
{
    auto event = static_cast<event_t *>(data);
    cout<<event->time_delta<<endl;
}

bool tfun::run(unsigned int timenum,string file)
{
    auto open_res = this->bpf->open_perf_buffer(BPF_OUTPUT_NAME,this->print_event);
    streambuf* coutbuf=cout.rdbuf();
    ofstream fout;
    if (open_res.code() != 0){
        cout << "错误！open_perf_buffer 执行失败" << endl;
        return false;
    }
    TerSignal::init();
    cout << "开始执行，按Ctrl^C停止" << endl;
    if (file!=""){
        fout.open(file);
        if (!fout.is_open())
            return false;
        cout.rdbuf(fout.rdbuf());
    }
    cout << "time_ns" << endl;
    if (timenum!=0)
        TerSignal::salarm(timenum);
    while (TerSignal::ok()){
        this->bpf->poll_perf_buffer(BPF_OUTPUT_NAME);
    }
    cout<<endl;
    if (file!=""){
        fout.flush();
        fout.close();
        cout.rdbuf(coutbuf);
    }
    cout<<"终止"<<endl;
    return true;
}